package com.star.langchain.learn.service;

import dev.langchain4j.agent.tool.ToolExecutionRequest;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * 工具调用确认服务
 * 用于管理待确认的工具调用请求，支持用户确认/拒绝和超时处理
 *
 * @author star
 */
@Slf4j
@Service
public class ToolConfirmationService {

    /**
     * 默认超时时间（秒）
     */
    private static final long DEFAULT_TIMEOUT_SECONDS = 30;

    /**
     * 存储待确认的工具调用请求
     * Key: 确认ID, Value: 确认结果Future
     */
    private final Map<String, ToolConfirmation> pendingConfirmations = new ConcurrentHashMap<>();

    public final ThreadLocal<String> confirmationIdThreadLocal = new ThreadLocal<>();


    /**
     * 工具确认请求结果
     */
    @Data
    public static class ConfirmationRequest {
        private String confirmationId;
        private CompletableFuture<Boolean> future;
    }

    public String getConfirmationId(){
        return confirmationIdThreadLocal.get();
    }
    public void clearConfirmationId(){
        confirmationIdThreadLocal.remove();
    }

    /**
     * 请求工具调用确认
     *
     * @param toolRequest 工具执行请求
     * @param timeoutSeconds 超时时间（秒）
     * @return 确认请求结果，包含确认ID和Future
     */
    public ConfirmationRequest requestConfirmation(ToolExecutionRequest toolRequest, long timeoutSeconds) {
        String confirmationId = UUID.randomUUID().toString();
        CompletableFuture<Boolean> future = new CompletableFuture<>();

        ToolConfirmation confirmation = new ToolConfirmation();
        confirmation.setConfirmationId(confirmationId);
        confirmation.setToolRequest(toolRequest);
        confirmation.setFuture(future);
        confirmation.setRequestTime(LocalDateTime.now());
        confirmation.setTimeoutSeconds(timeoutSeconds);

        confirmationIdThreadLocal.set(confirmationId);
        pendingConfirmations.put(confirmationId, confirmation);

        log.info("工具调用确认请求已创建: confirmationId={}, toolName={}, timeout={}秒", 
                confirmationId, toolRequest.name(), timeoutSeconds);

        // 设置超时处理
        future.orTimeout(timeoutSeconds, TimeUnit.SECONDS)
                .exceptionally(throwable -> {
                    if (throwable instanceof TimeoutException) {
                        log.warn("工具调用确认超时: confirmationId={}, toolName={}", 
                                confirmationId, toolRequest.name());
                        pendingConfirmations.remove(confirmationId);
                        clearConfirmationId();
                        return false; // 超时默认为拒绝
                    }
                    return false;
                });

        ConfirmationRequest request = new ConfirmationRequest();
        request.setConfirmationId(confirmationId);
        request.setFuture(future);
        return request;
    }

    /**
     * 请求工具调用确认（使用默认超时时间）
     *
     * @param toolRequest 工具执行请求
     * @return 确认请求结果，包含确认ID和Future
     */
    public ConfirmationRequest requestConfirmation(ToolExecutionRequest toolRequest) {
        return requestConfirmation(toolRequest, DEFAULT_TIMEOUT_SECONDS);
    }

    /**
     * 用户确认工具调用
     *
     * @param confirmationId 确认ID
     * @return 是否成功处理
     */
    public boolean confirm(String confirmationId) {
        ToolConfirmation confirmation = pendingConfirmations.remove(confirmationId);
        clearConfirmationId();
        if (confirmation != null) {
            log.info("用户确认工具调用: confirmationId={}, toolName={}", 
                    confirmationId, confirmation.getToolRequest().name());
            confirmation.getFuture().complete(true);
            return true;
        }
        log.warn("未找到待确认的工具调用: confirmationId={}", confirmationId);
        return false;
    }

    /**
     * 用户拒绝工具调用
     *
     * @param confirmationId 确认ID
     * @return 是否成功处理
     */
    public boolean reject(String confirmationId) {
        ToolConfirmation confirmation = pendingConfirmations.remove(confirmationId);
        clearConfirmationId();
        if (confirmation != null) {
            log.info("用户拒绝工具调用: confirmationId={}, toolName={}", 
                    confirmationId, confirmation.getToolRequest().name());
            confirmation.getFuture().complete(false);
            return true;
        }
        log.warn("未找到待确认的工具调用: confirmationId={}", confirmationId);
        return false;
    }

    /**
     * 获取待确认的工具调用信息
     *
     * @param confirmationId 确认ID
     * @return 工具调用信息，如果不存在则返回null
     */
    public ToolConfirmationInfo getConfirmationInfo(String confirmationId) {
        ToolConfirmation confirmation = pendingConfirmations.get(confirmationId);
        if (confirmation != null) {
            ToolConfirmationInfo info = new ToolConfirmationInfo();
            info.setConfirmationId(confirmationId);
            info.setToolName(confirmation.getToolRequest().name());
            info.setToolArguments(confirmation.getToolRequest().arguments());
            info.setRequestTime(confirmation.getRequestTime());
            info.setTimeoutSeconds(confirmation.getTimeoutSeconds());
            return info;
        }
        return null;
    }

    /**
     * 工具确认信息
     */
    @Data
    private static class ToolConfirmation {
        private String confirmationId;
        private ToolExecutionRequest toolRequest;
        private CompletableFuture<Boolean> future;
        private LocalDateTime requestTime;
        private long timeoutSeconds;
    }

    /**
     * 工具确认信息（对外暴露）
     */
    @Data
    public static class ToolConfirmationInfo {
        private String confirmationId;
        private String toolName;
        private String toolArguments;
        private LocalDateTime requestTime;
        private long timeoutSeconds;
    }
}

